from transformers import BertTokenizer, BertForSequenceClassification
import torch
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW
import torch.nn as nn


# 自定义数据集类
class SentimentDataset(Dataset):
  def __init__(self, texts, labels, tokenizer, max_length=128):
    self.texts = texts
    self.labels = labels
    self.tokenizer = tokenizer
    self.max_length = max_length


  def __len__(self):
    return len(self.texts)


  def __getitem__(self, idx):
    text = self.texts[idx]
    label = self.labels[idx]
    encoding = self.tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=self.max_length)
    input_ids = encoding['input_ids'].squeeze()
    attention_mask = encoding['attention_mask'].squeeze()
    return {
      'input_ids': input_ids,
      'attention_mask': attention_mask,
      'labels': torch.tensor(label)
    }


# 示例数据
texts = ["这部电影太棒了，我很喜欢。", "这个产品质量太差了，不推荐。"]
labels = [1, 0]  # 1 表示正面情感，0 表示负面情感


# 加载分词器和预训练模型
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=2)


# 数据集和数据加载器
dataset = SentimentDataset(texts, labels, tokenizer)
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)


# 定义优化器和损失函数
optimizer = AdamW(model.parameters(), lr=2e-5)
criterion = nn.CrossEntropyLoss()


# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()
for epoch in range(3):  # 假设训练 3 个 epoch
  for batch in dataloader:
    input_ids = batch['input_ids'].to(device)
    attention_mask = batch['attention_mask'].to(device)
    labels = batch['labels'].to(device)


    optimizer.zero_grad()
    outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
    loss = outputs.loss
    loss.backward()
    optimizer.step()


    print(f"Epoch {epoch + 1}, Loss: {loss.item()}")


# 验证模型（这里仅为简单示例）
model.eval()
with torch.no_grad():
  for batch in dataloader:
    input_ids = batch['input_ids'].to(device)
    attention_mask = batch['attention_mask'].to(device)
    labels = batch['labels'].to(device)


    outputs = model(input_ids, attention_mask=attention_mask)
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=1)


    correct = (predictions == labels).sum().item()
    total = labels.size(0)
    accuracy = correct / total
    print(f"Accuracy: {accuracy}")


# 保存微调后的模型
model.save_pretrained('fine_tuned_bert')
tokenizer.save_pretrained('fine_tuned_bert')